#include "Socket.hpp"
#include "Epoll.hpp"
#include "Util.hpp"
#include "Protocol.hpp"
#include <string>
#include <functional>
#include <unordered_map>
#include <cassert>
class Server;
class Connection
{
    typedef function<void(Connection *)> func_t;

public:
    Connection(int fd, Server *tsp) : _fd(fd), _tsp(tsp)
    {
    }
    ~Connection()
    {
        close(_fd);
    }
    void Register(func_t read, func_t write, func_t except)
    {
        _read = read;
        _write = write;
        _except = except;
    }
    void Close()
    {
        close(_fd);
    }

public:
    int _fd;
    string _inbuffer;
    string _outbuffer;
    func_t _read;
    func_t _write;
    func_t _except;
    Server *_tsp;
};
class Server
{
    typedef function<void(Connection *)> func_t;
    const int num = 128;

private:
    void addConnection(int fd, func_t read, func_t write, func_t except)
    {
        Connection *conn = new Connection(fd, this);
        conn->Register(read, write, except);
        _conns.insert(make_pair(fd, conn));
    }
    void Read(Connection *conn)
    {
        while (true)
        {
            char buffer[1024];
            ssize_t s = recv(conn->_fd, buffer, sizeof(buffer) - 1, 0);
            if (s > 0)
            {
                buffer[s] = 0;
                conn->_inbuffer += buffer;
                _server(conn);
            }
            else if (s == 0)
            {
                logMessage(NORMAL, "client quit");
                if (conn->_except)
                    conn->_except(conn);
                break;
            }
            else
            {
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                    break;
                else if (errno == EINTR)
                    continue;
                else
                {
                    if (conn->_except)
                        conn->_except(conn);
                    return;
                }
            }
        }
    }
    void Write(Connection *conn)
    {
        while (true)
        {
            ssize_t s = send(conn->_fd, conn->_outbuffer.c_str(), conn->_outbuffer.size(), 0);
            if (s >= 0)
            {
                if (!conn->_outbuffer.empty())
                    conn->_outbuffer.erase(0, s);
                else
                    break;
            }
            else
            {
                if (errno == EAGAIN || errno == EWOULDBLOCK)
                    break;
                else if (errno == EINTR)
                    continue;
                else
                {
                    if (conn->_except)
                        conn->_except(conn);
                    return;
                }
            }
        }
        uint32_t event = EPOLLIN | EPOLLET;
        if (!conn->_outbuffer.empty())
            event |= EPOLLOUT;
        _epoll.Control(conn->_fd, EPOLL_CTL_MOD, event);
    }
    void Except(Connection *conn)
    {
        _epoll.Control(conn->_fd, EPOLL_CTL_DEL, 0);
        _conns.erase(conn->_fd);
        conn->Close();
        logMessage(NORMAL, "fd close : [%d]", conn->_fd);
        delete conn;
    }
    void Accepter(Connection *conn)
    {
        while (true)
        {
            int err;
            int sock = _sock.Accept(err);
            if (sock > 0)
            {
                addConnection(sock,
                              std::bind(&Server::Read, this, std::placeholders::_1),
                              std::bind(&Server::Write, this, std::placeholders ::_1),
                              std::bind(&Server::Except, this, std::placeholders ::_1));
                Util::setNonBlock(sock);
                _epoll.Control(sock, EPOLL_CTL_ADD, EPOLLIN | EPOLLET);
            }
            else
            {
                if (err == EAGAIN || err == EWOULDBLOCK)
                    break;
                else if (err == EINTR)
                    continue;
                else
                    break;
            }
        }
    }
    bool isInConns(Connection *conn)
    {
        auto it = _conns.find(conn->_fd);
        return it != _conns.end();
    }

public:
    Server(const uint16_t port, func_t server) : _port(port), _num(num), _events(nullptr), _server(server)
    {
    }

    void init()
    {
        _sock.Socket();
        _sock.Bind(_port);
        _sock.Listen();
        _epoll.Create();
        Util::setNonBlock(_sock.getFd());
        bool ret = _epoll.Control(_sock.getFd(), EPOLL_CTL_ADD, EPOLLIN | EPOLLET);
        assert(ret);
        addConnection(_sock.getFd(), std::bind(&Server::Accepter, this, std::placeholders::_1), nullptr, nullptr);
        _events = new struct epoll_event[_num];
    }
    void start()
    {
        while (true)
        {
            int count = _epoll.Wait(_num, _events);
            for (int i = 0; i < count; ++i)
            {
                int fd = _events[i].data.fd;
                uint32_t event = _events[i].events;
                Connection *conn = _conns[fd];
                if (event & EPOLLERR)
                    event |= (EPOLLIN | EPOLLOUT);
                if (event & EPOLLHUP)
                    event |= (EPOLLIN | EPOLLOUT);
                if ((event & EPOLLIN) && conn->_read && isInConns(conn))
                {
                    conn->_read(conn);
                }
                if ((event & EPOLLOUT) && conn->_write && isInConns(conn))
                {
                    conn->_write(conn);
                }
            }
        }
    }

private:
    uint16_t _port;
    Sock _sock;
    Epoll _epoll;
    unordered_map<int, Connection *> _conns;
    int _num;
    struct epoll_event *_events;
    func_t _server;
};